from tensorflow.keras import Model
from tensorflow.keras.layers import Flatten,Conv2D,MaxPooling2D,Dense,Dropout, BatchNormalization

"""
卷积操作tf.keras.layers.Conv2D
用到参数：input_shape:输入数据的维度大小 如3通道64*64的输入图片 要设置为（64，64，3），
参数input_shape只需要在第一个网络层进行设置
-filters:表示卷积核的个数，可自己定义，输入为整数，其值等于卷积后输出的通道数
-kernel_size:卷积核的大小 通常是正整数或列表 例如kernel_size=3 kernel_size（3，3）都表示宽高大小为3
-strides:卷积核的滑动步长，通常是一个正整数或列表，列表可以包含2个或4个元素，默认为（1，1）
 例子：strides=2,strides=(2,2)都表示宽高方向的滑动步长为2
-padding:补0设置 可选项为"VALID"，"SAME"
  -如果padding="VALID" 不做补0
  -如果padding="SAME",程序会自动补0，使得输出feature map的宽高=ceil(输入图片的宽高/strides)
  ceil表示向上取整 当strides=1 输入输出宽高一致
- activation:激活函数
  可选项：
  -'sigmoid'  'tanh'  'relu'  'elu'或tf.keras.activations.elu(alpha=1.0)
   'selu'等激活函数  'swish'（tf.2.2以上才有） softmax激活函数  
- input shape 如果是第一层卷积，需要输入图片大小(height,width,channel)，比如input_shape(128,128,3)
- name 输入字符串，给该层设置的一个名称

池化操作 tf.keras.layers.Maxpool2D
用到的参数:
-pool_size：池化卷积核的大小，具体设置与tf.keras.layers.Conv2D的kernel_size一样
-strides：池化的滑动步长，与tf.keras.layers.Conv2D的strides一样
-name  输入字符串，给该层设置的一个名称

全连接操作 tf.keras.layers.Dense
用到的参数:
-units:输入整数 全连接层神经元个数
-activation:激活函数 分类网络的输出层 一般用softmax激活函数
-name  输入字符串，给该层设置的一个名称

展平操作tf.keras.layers.Flatten
举例说明：
[[1,2,3],[4,5,6]]-> [1,2,3,4,5,6]
"""

def get_norm(norm_type):
    if norm_type == "batch":
        return BatchNormalization()
    else:
        raise ValueError(f"Unrecognized norm_type {norm_type}")


#定义一个子类来搭建模型
class ConvModel(Model):
    def __init__(self):
        #父类初始化操作
        super(ConvModel,self).__init__()
        #卷积层conv_1_1  input_shape从(64,64,3)->(150,150,3)
        self.conv_1_1 = Conv2D(input_shape=(150,150,3),filters=32,kernel_size=3,
                               padding='same',activation='relu',name='conv_1_1')
        self.norm_1_1 = get_norm(norm_type='batch')  #添加bn层

        #卷积层conv_1_2
        self.conv_1_2 = Conv2D(filters=32,kernel_size=3,padding='same',
                               activation='relu',name='conv_1_2')
        self.norm_1_2 = get_norm(norm_type='batch')

        #池化层max_pool_1
        self.max_pool_1 = MaxPooling2D(pool_size=2,name='max_pool_1')

        #卷积层conv_2_1
        self.conv_2_1 = Conv2D(filters=64,kernel_size=3,padding='same',
                               activation='relu',name='conv_2_1')
        self.norm_2_1 = get_norm(norm_type='batch')

        # 卷积层conv_2_2
        self.conv_2_2 = Conv2D(filters=64, kernel_size=3, padding='same',
                              activation='relu', name='conv_2_2')
        self.norm_2_2 = get_norm(norm_type='batch')

        # 池化层max_pool_2
        self.max_pool_2 = MaxPooling2D(pool_size=2, name='max_pool_2')

        # 卷积层conv_3_1
        self.conv_3_1 = Conv2D(filters=128, kernel_size=3, padding='same',
                               activation='relu', name='conv_3_1')
        self.norm_3_1 = get_norm(norm_type='batch')

        # 卷积层conv_3_2
        self.conv_3_2 = Conv2D(filters=128, kernel_size=3, padding='same',
                               activation='relu', name='conv_3_2')
        self.norm_3_2 = get_norm(norm_type='batch')

        # 池化层max_pool_3
        self.max_pool_3 = MaxPooling2D(pool_size=2, name='max_pool_3')

        #展平层flatten
        self.flatten = Flatten(name='flatten')
        self.dropout = Dropout(0.5)
        self.flatten_norm = get_norm(norm_type='batch')

        #全连接层 softmax实现kaggle分类
        self.dense = Dense(units=10,activation="softmax",name='logit')

    def call(self,x):
        x = self.conv_1_1(x)
        x = self.norm_1_1(x)
        x = self.conv_1_2(x)
        x = self.norm_1_2(x)
        x = self.max_pool_1(x)
        x = self.conv_2_1(x)
        x = self.norm_2_1(x)
        x = self.conv_2_2(x)
        x = self.norm_2_2(x)
        x = self.max_pool_2(x)
        x = self.conv_3_1(x)
        x = self.norm_3_1(x)
        x = self.conv_3_2(x)
        x = self.norm_3_2(x)
        x = self.max_pool_3(x)
        x = self.flatten(x)
        x = self.dropout(x)
        x = self.flatten_norm(x)
        x = self.dense(x)
        return x
